#pragma once

#include "HookApi.h"
#include "SocketHelper.h"
#include <WinSock2.h>
#include <unordered_map>

typedef int (WINAPI *__pfnconnect)(SOCKET s, const struct sockaddr FAR * name, 
	int namelen);
typedef int (WINAPI *__pfnsend)(SOCKET s, const char FAR *buf, int len, int flags);
typedef int (WINAPI *__pfnsendto)(SOCKET s, const char* buf, int len, int flags, 
	const struct sockaddr* to, int tolen);
typedef int (WINAPI *__pfnrecv)(SOCKET s, char FAR *buf, int len, int flags);
typedef int (WINAPI *__pfnrecvfrom)(SOCKET s, char* buf, int len, int flags, 
	struct sockaddr* from, int* fromlen);
typedef int (WINAPI *__pfnWSASend)(SOCKET s, __in_ecount(dwBufferCount) LPWSABUF lpBuffers, 
	DWORD dwBufferCount, LPDWORD lpNumberOfBytesSent, DWORD dwFlags, LPWSAOVERLAPPED lpOverlapped, 
	LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine);
typedef int (WINAPI *__pfnWSARecv)(
	SOCKET s, LPWSABUF lpBuffers, DWORD dwBufferCount, LPDWORD lpNumberOfBytesRecvd,
	LPDWORD lpFlags, LPWSAOVERLAPPED lpOverlapped, LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine);
typedef int (WINAPI *__pfnWSARecvFrom)(
	SOCKET s, LPWSABUF lpBuffers, DWORD dwBufferCount,
	LPDWORD lpNumberOfBytesRecvd, LPDWORD lpFlags, struct sockaddr FAR *lpFrom,                            
	LPINT lpFromlen, LPWSAOVERLAPPED lpOverlapped, LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine);
typedef int (WINAPI *__pfnWSAGetOverlappedResult)(SOCKET s, LPWSAOVERLAPPED lpOverlapped, LPDWORD lpcbTransfer, 
	BOOL fWait, LPDWORD lpdwFlags);
typedef BOOL (WINAPI *__pfnCreateProcessAsUserW)(HANDLE hToken,
	LPCTSTR lpApplicationName,
	LPTSTR lpCommandLine,
	LPSECURITY_ATTRIBUTES lpProcessAttributes,
	LPSECURITY_ATTRIBUTES lpThreadAttributes,
	BOOL bInheritHandles,
	DWORD dwCreationFlags,
	LPVOID lpEnvironment,
	LPCTSTR lpCurrentDirectory,
	LPSTARTUPINFO lpStartupInfo,
	LPPROCESS_INFORMATION lpProcessInformation);

class CWinsockHook
{
public:
	CWinsockHook(void);
	~CWinsockHook(void);

public:
	static INT WINAPI hook_connect(DWORD RetAddr, __pfnconnect pfnconnect, 
		SOCKET s, const struct sockaddr FAR * name, int namelen);
	static INT WINAPI hook_send(DWORD RetAddr, __pfnsend pfnsend, 
		SOCKET s, const char FAR *buf, int len, int flags);
	static INT WINAPI hook_sendto(DWORD RetAddr, __pfnsendto pfnsendto, 
		SOCKET s, const char* buf, int len, int flags, const struct sockaddr* to, 
		int tolen);
	static INT WINAPI hook_recv(DWORD RetAddr, __pfnrecv pfnrecv, 
		SOCKET s, char FAR *buf, int len, int flags);
	static INT WINAPI hook_recvfrom(DWORD RetAddr, __pfnrecvfrom pfnrecvfrom, 
		SOCKET s, char* buf, int len, int flags, struct sockaddr* from, 
		int* fromlen);
	static INT WINAPI hook_WSASend(DWORD RetAddr, __pfnWSASend pfnWSASend,
		SOCKET s, __in_ecount(dwBufferCount) LPWSABUF lpBuffers, 
		DWORD dwBufferCount, LPDWORD lpNumberOfBytesSent, DWORD dwFlags, 
		LPWSAOVERLAPPED lpOverlapped, 
		LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine);
	static INT WINAPI hook_WSARecv(
		DWORD RetAddr, 
		__pfnWSARecv pfnWSARecv,
		SOCKET s,
		LPWSABUF lpBuffers,
		DWORD dwBufferCount,
		LPDWORD lpNumberOfBytesRecvd,
		LPDWORD lpFlags,
		LPWSAOVERLAPPED lpOverlapped,
		LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine
		);
	static INT WINAPI hook_WSAGetOverlappedResult(DWORD RetAddr, __pfnWSAGetOverlappedResult pfnWSAGetOverlappedResult, SOCKET s, LPWSAOVERLAPPED lpOverlapped, LPDWORD lpcbTransfer, BOOL fWait, LPDWORD lpdwFlags);
	static BOOL WINAPI hook_CreateProcessAsUserW(DWORD retAddr,
		__pfnCreateProcessAsUserW pfnCreateProcessAsUserW,
		HANDLE hToken,
		LPCTSTR lpApplicationName,
		LPTSTR lpCommandLine,
		LPSECURITY_ATTRIBUTES lpProcessAttributes,
		LPSECURITY_ATTRIBUTES lpThreadAttributes,
		BOOL bInheritHandles,
		DWORD dwCreationFlags,
		LPVOID lpEnvironment,
		LPCTSTR lpCurrentDirectory,
		LPSTARTUPINFO lpStartupInfo,
		LPPROCESS_INFORMATION lpProcessInformation);
	
	static std::unordered_map<SOCKET, CSocketHelper*> mapSocketHelper;

private:
	PHOOKENVIRONMENT pHookEnv_connect;
	PHOOKENVIRONMENT pHookEnv_send;
	PHOOKENVIRONMENT pHookEnv_sendto;
	PHOOKENVIRONMENT pHookEnv_recv;
	PHOOKENVIRONMENT pHookEnv_recvfrom;
	PHOOKENVIRONMENT pHookEnv_WSASend;
	PHOOKENVIRONMENT pHookEnv_WSARecv;
	PHOOKENVIRONMENT pHookEnv_WSARecvFrom;
	PHOOKENVIRONMENT pHookEnv_WSAGetOverlappedResult;
	PHOOKENVIRONMENT pHookEnv_CreateProcessAsUserW;

public:
	static VOID threadCacheStatusGuard(PVOID param);
};

